// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/core/expressions.h>
#include <pollux/common/encode/base64.h>
#include <pollux/vector/vector_saver.h>

namespace kumo::pollux::core {
    namespace {
        TypePtr deserializeType(const melon::Dynamic &obj, void *context) {
            return ISerializable::deserialize<Type>(obj["type"]);
        }

        std::vector<TypedExprPtr> deserializeInputs(
            const melon::Dynamic &obj,
            void *context) {
            if (obj.count("inputs")) {
                return ISerializable::deserialize<std::vector<ITypedExpr> >(
                    obj["inputs"], context);
            }

            return {};
        }
    } // namespace

    melon::Dynamic ITypedExpr::serializeBase(std::string_view name) const {
        melon::Dynamic obj = melon::Dynamic::object;
        obj["name"] = name;
        obj["type"] = type_->serialize();

        if (!inputs_.empty()) {
            melon::Dynamic serializedInputs = melon::Dynamic::array;
            for (const auto &input: inputs_) {
                serializedInputs.push_back(input->serialize());
            }

            obj["inputs"] = serializedInputs;
        }

        return obj;
    }

    // static
    void ITypedExpr::registerSerDe() {
        auto &registry = DeserializationWithContextRegistryForSharedPtr();

        registry.Register("CallTypedExpr", core::CallTypedExpr::create);
        registry.Register("CastTypedExpr", core::CastTypedExpr::create);
        registry.Register("ConcatTypedExpr", core::ConcatTypedExpr::create);
        registry.Register("ConstantTypedExpr", core::ConstantTypedExpr::create);
        registry.Register("DereferenceTypedExpr", core::DereferenceTypedExpr::create);
        registry.Register("FieldAccessTypedExpr", core::FieldAccessTypedExpr::create);
        registry.Register("InputTypedExpr", core::InputTypedExpr::create);
        registry.Register("LambdaTypedExpr", core::LambdaTypedExpr::create);
    }

    void InputTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic InputTypedExpr::serialize() const {
        return ITypedExpr::serializeBase("InputTypedExpr");
    }

    // static
    TypedExprPtr InputTypedExpr::create(const melon::Dynamic &obj, void *context) {
        auto type = core::deserializeType(obj, context);

        return std::make_shared<InputTypedExpr>(std::move(type));
    }

    void ConstantTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic ConstantTypedExpr::serialize() const {
        auto obj = ITypedExpr::serializeBase("ConstantTypedExpr");
        if (valueVector_) {
            std::ostringstream out;
            saveVector(*valueVector_, out);
            auto serializedValue = out.str();
            obj["value_vector"] = encoding::Base64::encode(
                serializedValue.data(), serializedValue.size());
        } else {
            obj["value"] = value_.serialize();
        }

        return obj;
    }

    // static
    TypedExprPtr ConstantTypedExpr::create(
        const melon::Dynamic &obj,
        void *context) {
        auto type = core::deserializeType(obj, context);

        if (obj.count("value")) {
            auto value = variant::create(obj["value"]);
            return std::make_shared<ConstantTypedExpr>(std::move(type), value);
        }

        auto encodedData = obj["value_vector"].asString();
        auto serializedData = encoding::Base64::decode(encodedData);
        std::istringstream dataStream(serializedData);

        auto *pool = static_cast<memory::MemoryPool *>(context);

        return std::make_shared<ConstantTypedExpr>(restoreVector(dataStream, pool));
    }

    void CallTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic CallTypedExpr::serialize() const {
        auto obj = ITypedExpr::serializeBase("CallTypedExpr");
        obj["functionName"] = name_;
        return obj;
    }

    // static
    TypedExprPtr CallTypedExpr::create(const melon::Dynamic &obj, void *context) {
        auto type = core::deserializeType(obj, context);
        auto inputs = deserializeInputs(obj, context);

        return std::make_shared<CallTypedExpr>(
            std::move(type), std::move(inputs), obj["functionName"].asString());
    }

    void FieldAccessTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic FieldAccessTypedExpr::serialize() const {
        auto obj = ITypedExpr::serializeBase("FieldAccessTypedExpr");
        obj["fieldName"] = name_;
        return obj;
    }

    // static
    TypedExprPtr FieldAccessTypedExpr::create(
        const melon::Dynamic &obj,
        void *context) {
        auto type = core::deserializeType(obj, context);
        auto inputs = deserializeInputs(obj, context);
        POLLUX_CHECK_LE(inputs.size(), 1);

        auto name = obj["fieldName"].asString();

        if (inputs.empty()) {
            return std::make_shared<FieldAccessTypedExpr>(std::move(type), name);
        } else {
            return std::make_shared<FieldAccessTypedExpr>(
                std::move(type), std::move(inputs[0]), name);
        }
    }

    void DereferenceTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic DereferenceTypedExpr::serialize() const {
        auto obj = ITypedExpr::serializeBase("DereferenceTypedExpr");
        obj["fieldIndex"] = index_;
        return obj;
    }

    // static
    TypedExprPtr DereferenceTypedExpr::create(
        const melon::Dynamic &obj,
        void *context) {
        auto type = core::deserializeType(obj, context);
        auto inputs = deserializeInputs(obj, context);
        POLLUX_CHECK_EQ(inputs.size(), 1);

        uint32_t index = obj["fieldIndex"].asInt();

        return std::make_shared<DereferenceTypedExpr>(
            std::move(type), std::move(inputs[0]), index);
    }

    void ConcatTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic ConcatTypedExpr::serialize() const {
        return ITypedExpr::serializeBase("ConcatTypedExpr");
    }

    // static
    TypedExprPtr ConcatTypedExpr::create(const melon::Dynamic &obj, void *context) {
        auto type = core::deserializeType(obj, context);
        auto inputs = deserializeInputs(obj, context);

        return std::make_shared<ConcatTypedExpr>(
            type->as_row().names(), std::move(inputs));
    }

    void LambdaTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic LambdaTypedExpr::serialize() const {
        auto obj = ITypedExpr::serializeBase("LambdaTypedExpr");
        obj["signature"] = signature_->serialize();
        obj["body"] = body_->serialize();
        return obj;
    }

    // static
    TypedExprPtr LambdaTypedExpr::create(const melon::Dynamic &obj, void *context) {
        auto signature = ISerializable::deserialize<Type>(obj["signature"]);
        auto body = ISerializable::deserialize<ITypedExpr>(obj["body"], context);

        return std::make_shared<LambdaTypedExpr>(
            as_row_type(signature), std::move(body));
    }

    void CastTypedExpr::accept(
        const ITypedExprVisitor &visitor,
        ITypedExprVisitorContext &context) const {
        visitor.visit(*this, context);
    }

    melon::Dynamic CastTypedExpr::serialize() const {
        auto obj = ITypedExpr::serializeBase("CastTypedExpr");
        obj["nullOnFailure"] = nullOnFailure_;
        return obj;
    }

    // static
    TypedExprPtr CastTypedExpr::create(const melon::Dynamic &obj, void *context) {
        auto type = core::deserializeType(obj, context);
        auto inputs = deserializeInputs(obj, context);

        return std::make_shared<CastTypedExpr>(
            std::move(type), std::move(inputs), obj["nullOnFailure"].asBool());
    }
} // namespace kumo::pollux::core
